package main

import (
	"encoding/json"
	"github.com/PuerkitoBio/goquery"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

const baseUrl = "http://www.stats.gov.cn/sj/tjbz/tjyqhdmhcxhfdm/2022/"

var extraCookie = ""

func main() {
	areas := InitArea()
	WriteJson(areas)
}

func WriteJson(areas *[]Area) {
	jsonData, err := json.MarshalIndent(areas, "", "  ")
	if err != nil {
		LogErr(err)
	}
	err = os.WriteFile("areas.json", jsonData, 0644)
	if err != nil {
		LogError("文件输出失败")
		LogErr(err)
	}
}

func InitArea() *[]Area {
	document, _ := NewDocument(baseUrl)
	elements := document.Find("tr.provincetr td a")
	var areas = &[]Area{}
	elements.Each(func(i int, element *goquery.Selection) {
		href, _ := element.Attr("href")
		area := Area{
			Label:    element.Text(),
			Code:     strings.ReplaceAll(href, ".html", ""),
			Sort:     i + 1,
			Children: Children(href),
		}
		*areas = append(*areas, area)
	})
	return areas
}

func Children(href string) *[]Area {
	if len(strings.TrimSpace(href)) == 0 {
		LogError("空链接")
		return nil
	}
	log.Println("URL-----", href)
	document, err := NewDocument(baseUrl + href)
	if err != nil {
		LogErr(err)
		time.Sleep(5 * time.Minute)
		return Children(href)
	}
	elements := document.Find("body > table:nth-child(3) > tbody > tr:nth-child(1) > td > table > tbody > tr:nth-child(2) > td > table > tbody > tr > td > table > tbody")
	if elements == nil || elements.Size() == 0 {
		if strings.Contains(document.Text(), "Please enable JavaScript and refresh the page.") {
			LogError("请求成功,但节点为空: " + (baseUrl + href))
			time.Sleep(15 * time.Minute)
		} else {
			panic("html解析失败: " + href)
		}
		return Children(href)
	}
	var areas = &[]Area{}
	elements.Children().Each(func(i int, element *goquery.Selection) {
		td := element.Children()
		area := Area{
			Code: td.First().Text(),
			Sort: i + 1,
		}
		if td.First().Children().Size() == 0 {
			area.Label = td.Eq(td.Size() - 1).Text()
		} else {
			area.Label = td.Eq(1).Text()
			url, _ := td.Eq(1).Children().First().Attr("href")
			prefix := href[0 : strings.LastIndex(href, "/")+1]
			if !strings.HasPrefix(url, prefix) {
				url = prefix + url
			}
			area.Children = Children(url)
		}
		*areas = append(*areas, area)
	})
	return areas
}

func NewDocument(url string) (*goquery.Document, error) {
	client := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	// 添加请求头
	req.Header.Add("Host", "www.stats.gov.cn")
	req.Header.Add("Upgrade-Insecure-Requests", "1")
	req.Header.Add("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7")
	req.Header.Add("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36 Edg/115.0.1901.203")

	content, _ := os.ReadFile("./cookie.json")
	if err != nil {
		LogError("cookie文件读取错误")
	} else {
		var cookie map[string]string
		err = json.Unmarshal(content, &cookie)
		if err != nil {
			LogError("json转换错误")
			LogErr(err)
		} else {
			var cookieValue = ""
			for k, v := range cookie {
				cookieValue += k + "=" + v + ";"
			}
			req.Header.Add("Cookie", cookieValue+extraCookie)
		}
	}
	// 发送请求
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer func(Body io.ReadCloser) {
		err := Body.Close()
		if err != nil {
		}
	}(resp.Body)
	if resp.StatusCode == 302 {
		ck := resp.Header.Get("Set-Cookie")
		split := strings.Split(ck, ";")
		for i := range split {
			if strings.Contains(split[i], "wzws_cid") {
				extraCookie = split[i]
				break
			}
		}
		return NewDocument(url)
	}
	// 使用goquery.NewDocumentFromReader函数创建一个文档对象
	// 在请求 11/01/12/110112004.html 时不知什么原因document初始化错误，标签错乱
	doc, _ := goquery.NewDocumentFromReader(resp.Body)
	return doc, nil
}

func LogErr(err error) {
	log.Println("error: ", err)
}

func LogError(err string) {
	log.Println("error: ", err)
}
